from abc import ABC, abstractmethod
import torch
import torch.nn as nn
from typing import Dict, Any, Optional

class BaseModel(nn.Module, ABC):
    """Abstract base class for all models"""
    
    def __init__(self):
        super().__init__()
        
    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass of the model
        
        Args:
            x (torch.Tensor): Input tensor
            
        Returns:
            torch.Tensor: Output tensor
        """
        pass
    
    @abstractmethod
    def configure_optimizers(self) -> torch.optim.Optimizer:
        """Configure optimizer for the model
        
        Returns:
            torch.optim.Optimizer: Configured optimizer
        """
        pass
    
    def save_model(self, path: str,
                   additional_info: Optional[Dict[str, Any]] = None) -> None:
        """Save model state and additional information
        
        Args:
            path (str): Path to save the model
            additional_info (Optional[Dict[str, Any]]): Additional information to save
        """
        save_dict = {
            'model_state_dict': self.state_dict(),
            'model_config': self.get_config()
        }
        if additional_info:
            save_dict.update(additional_info)
        torch.save(save_dict, path)
    
    def load_model(self, path: str) -> Dict[str, Any]:
        """Load model state and additional information
        
        Args:
            path (str): Path to load the model from
            
        Returns:
            Dict[str, Any]: Loaded information
        """
        checkpoint = torch.load(path)
        self.load_state_dict(checkpoint['model_state_dict'])
        return checkpoint
    
    @abstractmethod
    def get_config(self) -> Dict[str, Any]:
        """Get model configuration
        
        Returns:
            Dict[str, Any]: Model configuration
        """
        pass

from .transformer import TransformerModel
from .lstm import LSTMModel
from .tcn import TCNModel

__all__ = ['BaseModel', 'TransformerModel', 'LSTMModel', 'TCNModel']